import numpy as np
import cv2

def point_cloud_from_depth_c2w(
    image,             # HxW (gray) 或 HxWx3 (BGR/RGB)，若不需要颜色可传 None
    depth,             # HxW 深度图；单位米或由 depth_scale 转成米
    K,                 # 3x3 相机内参矩阵
    dist,              # 畸变系数 (OpenCV 常规模型: k1,k2,p1,p2,k3[,k4,k5,k6...])；若无畸变可传 None 或全 0
    c2w,               # 3x4 或 4x4 相机到世界的齐次变换矩阵；满足 [Pw;1] = c2w @ [Pc;1]
    depth_scale=0.001,   # 深度缩放（例如深度以毫米存储时设为 0.001）
    undistort=True,    # 是否对像素点做去畸变再反投影
    valid_min=1e-6,    # 有效深度下限
    valid_max=2.0,  # 有效深度上限
    depth_is_along_ray=False,  # True 表示 depth 为沿视线距离；False 表示为相机坐标系 Z（常见）
    return_world=True           # True 返回世界系点云；False 返回相机系点云
):
    """
    返回:
      points: Nx3 点云（世界或相机坐标系，取决于 return_world）
      colors: NxC 颜色（与 image 通道一致；若 image 为 None 则为 None）
      pixels: Nx2 像素坐标 (x,y)，对应每个点的原图像素位置
    """
    # 深度准备
    depth = np.asarray(depth).astype(np.float32) * depth_scale
    H, W = depth.shape[:2]

    # 有效深度掩码
    mask = np.isfinite(depth) & (depth > valid_min) & (depth < valid_max)
    if not np.any(mask):
        return np.zeros((0, 3), np.float32), None, np.zeros((0, 2), np.int32)

    ys, xs = np.where(mask)
    zs = depth[ys, xs]  # 可能是 Z 或 视线距离

    # 内参
    K = np.asarray(K, dtype=np.float64)
    fx, fy = K[0, 0], K[1, 1]
    cx, cy = K[0, 2], K[1, 2]

    # 计算归一化相机坐标 (x_n, y_n)
    use_dist = undistort and (dist is not None) and (np.size(dist) > 0) and np.any(np.array(dist) != 0)
    if use_dist:
        pts_pix = np.stack([xs, ys], axis=1).astype(np.float32).reshape(-1, 1, 2)
        # undistortPoints 返回理想归一化坐标（焦距=1，主点=0）
        pts_norm = cv2.undistortPoints(pts_pix, K, np.asarray(dist, dtype=np.float64), R=None, P=None)
        x_n = pts_norm[:, 0, 0]
        y_n = pts_norm[:, 0, 1]
    else:
        x_n = (xs - cx) / fx
        y_n = (ys - cy) / fy

    # 归一化坐标 + 深度 -> 相机坐标系点
    if depth_is_along_ray:
        ray_norm = np.sqrt(x_n**2 + y_n**2 + 1.0)
        scale = zs / ray_norm
        Xc = x_n * scale
        Yc = y_n * scale
        Zc = 1.0 * scale
    else:
        Xc = x_n * zs
        Yc = y_n * zs
        Zc = zs

    Pc = np.stack([Xc, Yc, Zc], axis=1).astype(np.float64)  # N x 3

    # 相机系 -> 世界系（如需）
    if return_world:
        c2w = np.asarray(c2w, dtype=np.float64)
        if c2w.shape == (3, 4):
            c2w = np.vstack([c2w, [0, 0, 0, 1]])
        if c2w.shape != (4, 4):
            raise ValueError("c2w 必须是 4x4 或 3x4 的齐次变换矩阵")
        Pc_h = np.hstack([Pc, np.ones((Pc.shape[0], 1), dtype=np.float64)])   # N x 4
        Pw_h = (c2w @ Pc_h.T).T
        points = Pw_h[:, :3].astype(np.float32)
    else:
        points = Pc.astype(np.float32)

    # 采样颜色
    colors = None
    if image is not None:
        img = np.asarray(image, dtype=np.float32) / 255.
        if img.ndim == 2:
            colors = img[ys, xs][:, None]
        else:
            colors = img[ys, xs, :]

    pixels = np.stack([xs, ys], axis=1).astype(np.int32)
    return points, colors, pixels